#!/bin/bash
# 版本: 1.0
# 作者: huangsheng
# 兼容系统版本:ubuntu
# 验证:未验证

# 定义变量
software_dir="stable-diffusion"
venv_name="amd${software_dir}"
ONEGPU_RUN_SD="onegpu_run_sd.sh"
MGPU_INSTALL="mgpu_sd_install.sh"
SCRIPT_PATH=$(realpath "$0")

# 检查9090端口是否占用
DEFAULT_PORT=9090
while netstat -tuln | grep -q ":$DEFAULT_PORT"; do
    echo "端口 $DEFAULT_PORT 已被占用。请输入一个新的端口号："
    read -r DEFAULT_PORT
done
echo "使用的端口号: $DEFAULT_PORT"

# 从文件中提取HSA_OVERRIDE_GFX_VERSION的值
if [ -f "$ONEGPU_RUN_SD" ]; then
    found_version=$(grep 'export HSA_OVERRIDE_GFX_VERSION=' "$ONEGPU_RUN_SD" | awk -F'=' '{print $2}')
    echo "使用的HSA_OVERRIDE_GFX_VERSION: $found_version"
else
    echo "找不到$ONEGPU_RUN_SD文件。请确保路径正确。"
    exit 1
fi

# 获取本机上网IP地址
IP_ADDR=$(hostname -I | awk '{print $1}')

# 查询系统中的显卡数量
GPU_COUNT=$(lspci | grep -i 'vga compatible controller: Advanced Micro Devices' | wc -l)

# 确认获取到的GPU数量
if [ "$GPU_COUNT" -eq 0 ]; then
    echo "未检测到AMD显卡，脚本退出。"
    exit 1
fi

echo "检测到的AMD显卡数量: $GPU_COUNT"
echo "本机上网IP地址: $IP_ADDR"

# 移除旧的 Nginx 配置文件
sudo rm -rf /etc/nginx/nginx.conf

# 修改Nginx配置文件
sudo bash -c "cat > /etc/nginx/nginx.conf" <<EOL
events {
    worker_connections 1024;
}

http {
    include       mime.types;
    default_type  application/octet-stream;
    
    map \$http_upgrade \$connection_upgrade {
        default upgrade;
        ''      close;
    }

    sendfile        on;
    keepalive_timeout  65;

    upstream sd {
        hash \$remote_addr;
EOL

# 添加每个实例的后端服务器配置
for i in $(seq 0 $((GPU_COUNT - 1))); do
    port=$((7860 + i))
    sudo bash -c "echo '        server $IP_ADDR:$port;' >> /etc/nginx/nginx.conf"
done

sudo bash -c "cat >> /etc/nginx/nginx.conf" <<EOL
    }

    server {
        listen $DEFAULT_PORT;
        server_name $IP_ADDR;

        location / {
            proxy_pass http://sd;
            proxy_set_header Upgrade           \$http_upgrade;
            proxy_set_header Connection        \$connection_upgrade;
            proxy_set_header X-Real-IP         \$remote_addr;
            proxy_set_header X-Forwarded-For   \$proxy_add_x_forwarded_for;
        }
    }
}
EOL

# 生成启动脚本，名称已修改为mgpu_run_sd.sh
cat <<EOL > mgpu_run_sd.sh
#!/bin/bash

# 激活虚拟环境
source $venv_name/bin/activate
export HSA_OVERRIDE_GFX_VERSION=$found_version

EOL

# 添加每个实例的启动命令
for i in $(seq 0 $((GPU_COUNT - 1))); do
    port=$((7860 + i))
    cat <<EOL >> mgpu_run_sd.sh

# 启动第$((i+1))个实例
export HIP_VISIBLE_DEVICES=$i
python3 launch.py --port $port --precision full --no-half --autolaunch --listen --theme dark --skip-torch-cuda-test --server-name $IP_ADDR &
EOL
done

# 确认所有实例已启动
cat <<EOL >> mgpu_run_sd.sh

echo "所有实例已启动。"
EOL

# 赋予启动脚本执行权限
chmod +x mgpu_run_sd.sh

echo "启动脚本 mgpu_run_sd.sh 已生成。"

# 启动并配置Nginx服务为开机自启动
sudo systemctl restart nginx
sudo systemctl enable nginx

# 删除脚本文件自身
sudo rm -- "$SCRIPT_PATH"
echo "Nginx已启动并配置为开机自启动，监听端口：$DEFAULT_PORT。"

